import os
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import threading
from datetime import datetime
from queue import Queue


class DirectorySizeViewer:
    def __init__(self, root):
        self.root = root
        self.root.title("目录大小查看器")
        self.root.geometry("900x650")

        # 历史记录堆栈
        self.history = []
        self.history_index = -1

        # 创建UI
        self.create_widgets()

        # 初始化变量
        self.current_path = os.path.expanduser('~')
        self.update_current_path()

        # 绑定事件
        self.tree.bind('<Double-1>', self.on_double_click)
        self.tree.bind('<Button-3>', self.show_context_menu)

        # 右键菜单
        self.create_context_menu()

        # 初始加载
        self.load_directory()

    def create_widgets(self):
        # 工具栏
        toolbar = tk.Frame(self.root, bd=1, relief=tk.RAISED)

        # 导航按钮
        self.back_btn = tk.Button(
            toolbar, text="←", command=self.go_back, width=3)
        self.back_btn.pack(side=tk.LEFT, padx=2, pady=2)

        self.forward_btn = tk.Button(
            toolbar, text="→", command=self.go_forward, width=3)
        self.forward_btn.pack(side=tk.LEFT, padx=2, pady=2)

        self.up_btn = tk.Button(toolbar, text="↑", command=self.go_up, width=3)
        self.up_btn.pack(side=tk.LEFT, padx=2, pady=2)

        # 地址栏
        self.path_var = tk.StringVar()
        self.path_entry = tk.Entry(
            toolbar, textvariable=self.path_var, width=70)
        self.path_entry.pack(side=tk.LEFT, padx=5, pady=2,
                             fill=tk.X, expand=True)
        self.path_entry.bind('<Return>', self.on_path_entered)

        # 浏览按钮
        browse_btn = tk.Button(toolbar, text="浏览...",
                               command=self.browse_directory)
        browse_btn.pack(side=tk.LEFT, padx=2, pady=2)

        toolbar.pack(side=tk.TOP, fill=tk.X)

        # 主面板
        main_panel = tk.Frame(self.root)

        # 树形视图
        self.tree = ttk.Treeview(main_panel, columns=(
            'size', 'type', 'modified'), selectmode='browse')

        # 设置列
        self.tree.column('#0', width=400, anchor=tk.W)
        self.tree.column('size', width=120, anchor=tk.E)
        self.tree.column('type', width=100, anchor=tk.W)
        self.tree.column('modified', width=150, anchor=tk.W)

        # 设置列标题
        self.tree.heading(
            '#0', text='名称', command=lambda: self.sort_treeview(0, False, False))
        self.tree.heading('size', text='大小 ▼',
                          command=lambda: self.sort_treeview(1, True, True))
        self.tree.heading('type', text='类型',
                          command=lambda: self.sort_treeview(2, False, False))
        self.tree.heading('modified', text='修改日期',
                          command=lambda: self.sort_treeview(3, False, False))

        # 滚动条
        vsb = ttk.Scrollbar(main_panel, orient="vertical",
                            command=self.tree.yview)
        hsb = ttk.Scrollbar(main_panel, orient="horizontal",
                            command=self.tree.xview)
        self.tree.configure(yscrollcommand=vsb.set, xscrollcommand=hsb.set)

        # 布局
        self.tree.grid(row=0, column=0, sticky='nsew')
        vsb.grid(row=0, column=1, sticky='ns')
        hsb.grid(row=1, column=0, sticky='ew')

        main_panel.grid_rowconfigure(0, weight=1)
        main_panel.grid_columnconfigure(0, weight=1)
        main_panel.pack(fill=tk.BOTH, expand=True)

        # 状态栏
        self.status_var = tk.StringVar()
        self.status_bar = tk.Label(
            self.root, textvariable=self.status_var, bd=1, relief=tk.SUNKEN, anchor=tk.W)
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)

        # 进度条
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(
            self.root, variable=self.progress_var, maximum=100)

        # 排序状态
        self.sort_column = 1  # 默认按大小排序
        self.sort_descending = True

    def create_context_menu(self):
        self.context_menu = tk.Menu(self.root, tearoff=0)
        self.context_menu.add_command(label="打开", command=self.open_selected)
        self.context_menu.add_command(
            label="在资源管理器中显示", command=self.reveal_in_explorer)
        self.context_menu.add_separator()
        self.context_menu.add_command(label="复制路径", command=self.copy_path)
        self.context_menu.add_command(label="属性", command=self.show_properties)

    def update_current_path(self):
        self.path_var.set(self.current_path)
        self.root.title(f"目录大小查看器 - {self.current_path}")

    def browse_directory(self):
        directory = filedialog.askdirectory(initialdir=self.current_path)
        if directory:
            self.add_to_history()
            self.current_path = directory
            self.update_current_path()
            self.load_directory()

    def on_path_entered(self, event):
        path = self.path_var.get()
        if os.path.isdir(path):
            self.add_to_history()
            self.current_path = path
            self.update_current_path()
            self.load_directory()
        else:
            messagebox.showerror("错误", "指定的路径不存在或不是目录")

    def load_directory(self):
        # 清空当前树
        for item in self.tree.get_children():
            self.tree.delete(item)

        # 显示加载状态
        self.status_var.set("正在扫描目录...")
        self.progress_bar.pack(fill=tk.X)
        self.root.update()

        # 使用队列传递结果
        self.result_queue = Queue()

        # 启动扫描线程
        scan_thread = threading.Thread(
            target=self.scan_directory, args=(self.current_path,), daemon=True)
        scan_thread.start()

        # 启动UI更新线程
        self.root.after(100, self.update_ui_from_queue)

    def scan_directory(self, path):
        try:
            items = []
            entries = list(os.scandir(path))
            total = len(entries)

            for i, entry in enumerate(entries):
                try:
                    if entry.is_dir():
                        item_type = "文件夹"
                        size = self.get_folder_size(entry.path)
                    else:
                        item_type = "文件"
                        size = entry.stat().st_size

                    modified = datetime.fromtimestamp(
                        entry.stat().st_mtime).strftime('%Y-%m-%d %H:%M:%S')

                    items.append({
                        'name': entry.name,
                        'size': size,
                        'formatted_size': self.format_size(size),
                        'type': item_type,
                        'modified': modified,
                        'path': entry.path,
                        'is_dir': entry.is_dir()
                    })

                    # 更新进度
                    progress = (i + 1) / total * 100
                    self.result_queue.put(('progress', progress))
                except (PermissionError, OSError) as e:
                    continue

            # 发送结果
            self.result_queue.put(('items', items))

        except Exception as e:
            self.result_queue.put(('error', str(e)))

    def update_ui_from_queue(self):
        try:
            while not self.result_queue.empty():
                result_type, data = self.result_queue.get_nowait()

                if result_type == 'progress':
                    self.progress_var.set(data)
                    self.root.update()

                elif result_type == 'items':
                    self.display_items(data)
                    self.progress_bar.pack_forget()

                elif result_type == 'error':
                    messagebox.showerror("错误", f"扫描目录时出错: {data}")
                    self.status_var.set(f"错误: {data}")
                    self.progress_bar.pack_forget()

        except Queue.Empty:
            pass

        # 继续检查队列
        if not hasattr(self, 'scan_thread') or self.scan_thread.is_alive():
            self.root.after(100, self.update_ui_from_queue)

    def display_items(self, items):
        # 清空当前树
        for item in self.tree.get_children():
            self.tree.delete(item)

        # 添加新项目
        for item in items:
            self.tree.insert('', 'end',
                           text=item['name'],
                           values=(item['formatted_size'],
                                   item['type'],
                                   item['modified']),
                           tags=('dir' if item['is_dir'] else 'file',),
                           iid=item['path'])

        # 配置标签样式
        self.tree.tag_configure('dir', foreground='blue')
        self.tree.tag_configure('file', foreground='black')

        # 更新状态栏
        dir_count = sum(1 for item in items if item['is_dir'])
        file_count = len(items) - dir_count
        total_size = sum(item['size'] for item in items)
        self.status_var.set(
            f"项目: {len(items)} (文件夹: {dir_count}, 文件: {file_count}) | "
            f"总大小: {self.format_size(total_size)} | "
            f"排序方式: {self.get_sort_description()}"
        )

        # 应用当前排序
        self.apply_sort()

    def get_folder_size(self, path):
        """递归计算文件夹大小（优化版）"""
        total = 0
        try:
            with os.scandir(path) as it:
                for entry in it:
                    try:
                        if entry.is_file():
                            total += entry.stat().st_size
                        elif entry.is_dir():
                            total += self.get_folder_size(entry.path)
                    except (PermissionError, OSError):
                        continue
        except (PermissionError, OSError):
            pass
        return total

    def format_size(self, size):
        """格式化文件大小为易读格式"""
        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
            if size < 1024.0:
                return f"{size:.2f} {unit}"
            size /= 1024.0
        return f"{size:.2f} PB"

    def on_double_click(self, event):
        item = self.tree.selection()[0]
        item_values = self.tree.item(item, 'values')

        if item_values and item_values[1] == "文件夹":
            self.add_to_history()
            self.current_path = item
            self.update_current_path()
            self.load_directory()

    def add_to_history(self):
        """添加当前路径到历史记录"""
        if self.history_index < len(self.history) - 1:
            self.history = self.history[:self.history_index + 1]
        self.history.append(self.current_path)
        self.history_index = len(self.history) - 1
        self.update_nav_buttons()

    def go_back(self):
        if self.history_index > 0:
            self.history_index -= 1
            self.current_path = self.history[self.history_index]
            self.update_current_path()
            self.load_directory()
            self.update_nav_buttons()

    def go_forward(self):
        if self.history_index < len(self.history) - 1:
            self.history_index += 1
            self.current_path = self.history[self.history_index]
            self.update_current_path()
            self.load_directory()
            self.update_nav_buttons()

    def go_up(self):
        parent = os.path.dirname(self.current_path)
        if os.path.isdir(parent):
            self.add_to_history()
            self.current_path = parent
            self.update_current_path()
            self.load_directory()

    def update_nav_buttons(self):
        self.back_btn['state'] = 'normal' if self.history_index > 0 else 'disabled'
        self.forward_btn['state'] = 'normal' if self.history_index < len(
            self.history) - 1 else 'disabled'

    def show_context_menu(self, event):
        item = self.tree.identify_row(event.y)
        if item:
            self.tree.selection_set(item)
            self.context_menu.post(event.x_root, event.y_root)

    def open_selected(self):
        item = self.tree.selection()[0]
        item_values = self.tree.item(item, 'values')

        if item_values and item_values[1] == "文件夹":
            self.add_to_history()
            self.current_path = item
            self.update_current_path()
            self.load_directory()
        else:
            try:
                os.startfile(item)
            except Exception as e:
                messagebox.showerror("错误", f"无法打开文件: {str(e)}")

    def reveal_in_explorer(self):
        item = self.tree.selection()[0]
        try:
            if os.path.isfile(item):
                os.system(f'explorer /select,"{item}"')
            else:
                os.startfile(item)
        except Exception as e:
            messagebox.showerror("错误", f"无法在资源管理器中显示: {str(e)}")

    def copy_path(self):
        item = self.tree.selection()[0]
        self.root.clipboard_clear()
        self.root.clipboard_append(item)
        self.status_var.set(f"已复制路径: {item}")

    def show_properties(self):
        item = self.tree.selection()[0]
        try:
            stat = os.stat(item)
            size = stat.st_size if os.path.isfile(
                item) else self.get_folder_size(item)

            props = (
                f"路径: {item}\n"
                f"类型: {'文件夹' if os.path.isdir(item) else '文件'}\n"
                f"大小: {self.format_size(size)}\n"
                f"修改时间: {datetime.fromtimestamp(stat.st_mtime)}\n"
                f"创建时间: {datetime.fromtimestamp(stat.st_ctime)}"
            )

            messagebox.showinfo("属性", props)
        except Exception as e:
            messagebox.showerror("错误", f"无法获取属性: {str(e)}")

    def sort_treeview(self, col, is_numeric=False, is_date=False):
        """对树形视图进行排序"""
        # 如果点击的是当前排序列，则切换排序方向
        if col == self.sort_column:
            self.sort_descending = not self.sort_descending
        else:
            # 点击了新列，默认降序排列
            self.sort_column = col
            self.sort_descending = True
        
        # 更新所有列标题，清除旧箭头，添加新箭头
        columns = ['#0', 'size', 'type', 'modified']  # 对应 Treeview 的列标识符
        for i, col_id in enumerate(columns):
            heading_text = self.tree.heading(col_id)['text']
            # 移除旧的排序箭头
            heading_text = heading_text.replace(' ▼', '').replace(' ▲', '')
            # 如果是当前排序列，添加新箭头
            if i == col:
                heading_text += ' ▼' if self.sort_descending else ' ▲'
            self.tree.heading(col_id, text=heading_text)

    def apply_sort(self):
        """应用当前排序设置"""
        items = []
        for item in self.tree.get_children():
            values = self.tree.item(item, 'values')
            path = item
            
            # 获取原始数据
            try:
                if os.path.isdir(path):
                    size = self.get_folder_size(path)
                    item_type = "文件夹"
                else:
                    size = os.path.getsize(path)
                    item_type = "文件"
                modified = datetime.fromtimestamp(os.path.getmtime(path)).strftime('%Y-%m-%d %H:%M:%S')
            except:
                size = 0
                item_type = "未知"
                modified = "未知"
            
            items.append({
                'item': item,
                'name': self.tree.item(item, 'text'),
                'size': size,
                'type': item_type,
                'modified': modified,
                'path': path
            })
        
        # 确定排序键函数
        if self.sort_column == 0:  # 名称
            key_func = lambda x: x['name'].lower()
        elif self.sort_column == 1:  # 大小
            key_func = lambda x: x['size']
        elif self.sort_column == 2:  # 类型
            key_func = lambda x: x['type']
        elif self.sort_column == 3:  # 修改日期
            key_func = lambda x: x['modified']
        else:
            return  # 无效的列索引
        
        # 执行排序
        items.sort(key=key_func, reverse=self.sort_descending)
        
        # 重新插入项目
        for index, item_data in enumerate(items):
            self.tree.move(item_data['item'], '', index)
        
        # 更新状态栏排序信息
        self.update_sort_status()

    def update_sort_status(self):
        """更新状态栏中的排序信息"""
        columns = ['名称', '大小', '类型', '修改日期']
        direction = '降序' if self.sort_descending else '升序'
        sort_info = f"排序方式: {columns[self.sort_column]} ({direction})"
        
        # 保留状态栏前面的信息，只更新排序部分
        current_status = self.status_var.get()
        if '|' in current_status:
            new_status = current_status.split('|')[0] + f"| {sort_info}"
        else:
            new_status = f"{current_status} | {sort_info}"
        
        self.status_var.set(new_status)
        
    def get_sort_description(self):
        columns = ['名称', '大小', '类型', '修改日期']
        direction = '降序' if self.sort_descending else '升序'
        return f"{columns[self.sort_column]} ({direction})"

def main():
    root = tk.Tk()
    app = DirectorySizeViewer(root)
    root.mainloop()

if __name__ == "__main__":
    main()